import requests,json 
from enum import Enum
from support.utils.toolkit import Log

class Client():
    DEFAULT_SERVER_URL="http://127.0.0.1:11434"
    DEFAULT_EMBED_MODEL="xiaoming/bce-embedding-base_v1:latest"
    DEFAULT_LLM_MODEL="deepseek-r1:1.5b"
    DEFAULT_CODE_FMT = "utf-8"
    DEFAULT_STREAM_CHUNK = 8192
    DEFAULT_CHAT_CONTEXT_COUNT = 20
    DEFAULT_SHOW_THOUGHT_CHAIN = True
    API_ENDPOINT = Enum("API",{
        "EMBEDDING":"api/embed",
        "CHAT":"api/chat",
        "GENERATE":"api/generate",
        "PULL":"api/pull",
        "LIST":"api/tags",
        "SHOW":"api/show"
    })
    HTTP_CODE = Enum("HTTP",{
        "SUCCESS":200
    })

    def __init__(self, *, url=None, model=None):
        self.baseUrl = url if url else Client.DEFAULT_SERVER_URL
        self.model = model if model else Client.DEFAULT_LLM_MODEL
        self.messages = []
        self.MaxContextCount = Client.DEFAULT_CHAT_CONTEXT_COUNT
        
    def __post__(self, url, data):
        response = None
        if(data):
            response = requests.post(url, json=data)
        else:
            response = requests.post(url)
        return response.json()
    
    def __get__(self, url):
        response = requests.get(url)
        return response.json()
    
    def __poststream__(self, url, data):
        response = requests.post(url, json=data, stream=True)
        if response.status_code == Client.HTTP_CODE.SUCCESS.value:
            for chunk in response.iter_content(chunk_size=Client.DEFAULT_STREAM_CHUNK):
                if chunk:
                    yield json.loads(chunk.decode(Client.DEFAULT_CODE_FMT))

    def generate(self, prompt):
        url = f"{self.baseUrl}/{Client.API_ENDPOINT.GENERATE.value}"
        data = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "temperature": 0.8
        }
        return self.__post__(url, data)

    def __removeChatMessage__(self, index=-1):
        chatcount = int((len(self.messages) - 1)/2)
        if(index < 0 and abs(index) < chatcount):
            index = (index + chatcount) * 2
        elif(index >= 0 and index < chatcount):
            index = index * 2
        else:
            index = (chatcount - 1) * 2
        index += 1
        last = self.messages[index]
        lastSecond = self.messages[index + 1]
        if(last["role"] == "assistant" and lastSecond["role"] == "user"):
            del self.messages[index:index+1]
            
    def chatWith(self, role, message, *, stream=True, new=False, replaceLast=False):
        if(new):
            self.messages.clear()
            self.messages.append({"role": "system", "content": role.ROLE_DEF})
        if(len(self.messages) > 3 and replaceLast):
            self.__removeChatMessage__()
        if(len(self.messages) > self.MaxContextCount):
            self.__removeChatMessage__(0)
        self.messages.append({"role": "user", "content": f"{message}"})
        url = f"{self.baseUrl}/{Client.API_ENDPOINT.CHAT.value}"
        data = {
            "model": self.model,
            "messages": self.messages,
            "stream": stream,
            "temperature": 0.8,
            "max_tokens": 150
        }
        response = ""
        thinkchain = False
        for chunk in self.__poststream__(url, data):
            phrase = chunk["message"]["content"]
            role = chunk["message"]["role"]
            status = Log.STREAM.KEEP
            if(phrase == "<think>"):
                thinkchain = True
                status = Log.STREAM.START
            elif(phrase == "</think>"):
                status = Log.STREAM.CHANGE
                thinkchain = False
            elif(not thinkchain):
                response += phrase
            if(chunk['done'] in ['true', 'True', True]):
                status = Log.STREAM.STOP
                self.messages.append({"role": chunk["message"]["role"], "content": f"{response}"})
            if(Client.DEFAULT_SHOW_THOUGHT_CHAIN or not thinkchain):
                Log.chat(role, phrase, thinkchain, status)
        return response.strip()

    def info(self, model):
        url = f"{self.baseUrl}/{Client.API_ENDPOINT.SHOW.value}"
        data = {
            "model": f"{model}",
            "verbose": True
        }
        return self.__post__(url, data)
        
    def list(self):
        url = f"{self.baseUrl}/{Client.API_ENDPOINT.LIST.value}"
        return self.__get__(url)

